#include "util.h"
#include "log.h"
#include "conf.h"
#include "main.h"

static int send_handshake(tcp_connection* conn)
{
	struct sockaddr_in dest_addr = { 0 };
	
	if (get_dest_addr(conn->local_bev, &dest_addr))
	{
		log_error("get_dest_addr() failed.");
		return -1;
	}
	
	if (bufferevent_write(conn->remote_bev, &dest_addr, sizeof(dest_addr)))
	{
		log_error("write handshake message failed.");
		return -1;
	}

	log_info("write handshake message successfully.");
	return 0;
}

static void data_read_callback(struct bufferevent* bev, void* ctx)
{
	tcp_connection* conn = (tcp_connection*)ctx;
	struct bufferevent* peer = get_peer_bufferevent(bev, conn);

	size_t bytes = relay_data(bev, peer);
	if (bytes > 0)
		enable_relay(conn);
	else
		destroy_connection(conn);
}

static void connection_event_callback(struct bufferevent* bev, short what, void* ctx)
{
	tcp_connection* conn = (tcp_connection*)ctx;
	log_connection_event(bev, what);
	int err = 1;

	if (what == BEV_EVENT_CONNECTED)
	{
		assert(bev == conn->remote_bev);
		if (send_handshake(conn) == 0)
		{
			err = 0;
			enable_relay(conn);
		}
	}

	if (err)
		destroy_connection(conn);
}

static void accept_callback(struct evconnlistener* listener, evutil_socket_t fd, struct sockaddr* addr, int addr_len, void* ctx)
{
	struct event_base* evbase = evconnlistener_get_base(listener);
	tcp_connection* conn = create_connection(0);

	struct sockaddr_in server_addr = { 0 };
	server_addr.sin_family = AF_INET;
	server_addr.sin_addr.s_addr = inet_addr(get_config()->server_addr);
	server_addr.sin_port = htons(get_config()->server_port);	

	do
	{
		conn->local_bev = tcp_accept(evbase, fd, data_read_callback, NULL, connection_event_callback, conn);
		if (!conn->local_bev)
			break;

		conn->remote_bev = tcp_connect(&server_addr, evbase, data_read_callback, NULL, connection_event_callback, conn);
		if (!conn->remote_bev)
			break;

		return;
	} while (0);

	evutil_closesocket(fd);
	destroy_connection(conn);
}

int main(int argc, char** argv)
{
	return generic_main(false, argc, argv, accept_callback, NULL, NULL);
}
